from draw_FFT import draw_FFT
from tqdm import tqdm
from matplotlib import pyplot as plt
import numpy as np


def cal_A_n(amp_base, i):  # 计算指数衰减的得到的振幅
    return amp_base * (0.4 ** (i - 1))


def filter_harmonic(Path):  # 过滤谐波
    freq_axis, amp_last, sample_rate = draw_FFT(Path)
    amp_last = list(amp_last)
    freq_axis = list(freq_axis)
    freq_axis = [int(i) for i in freq_axis]  # 为了方便计算，全部化为整数
    amp_new = [0] * len(amp_last)

    print("----------开始过滤谐波----------")
    for i in tqdm(range(0, len(freq_axis), 600)):  # 步长设置为100进行提取频率
        f = freq_axis[i]  # 获得基准频率f
        amp_base = amp_last[i]  # 获得f处的基准实际振幅
        amp_harm = 0  # 记录谐波的量

        amp_list = []  # 记录对应的频率
        for j in range(1, 11):
            if j * f in freq_axis:
                amp_list.append((j, amp_last[freq_axis.index(j * f)]))  # 获得10个或者更少频率对应的(第i位，i位的振幅)
            else:
                amp_list.append((j, None))

        for mul in range(1, 11):  # 以间隔为0.1的乘积因子进行选取最大概率的谐波含量
            for tur in amp_list:  # 对<=10的频率倍数振幅进行遍历
                amp_base_harm = amp_base * (mul / 10)
                idx, amp_i_fact = tur
                if amp_i_fact is None:
                    continue
                else:
                    amp_i_ideal = cal_A_n(amp_base_harm, idx)  # 计算i位置的指数衰减值
                    if abs(amp_i_ideal - amp_i_fact) < 100:
                        amp_harm = amp_base_harm  # 确定下来最大概率的谐波振幅
                        break
                    else:
                        continue

        amp_new[i] = amp_last[i] - amp_harm  # 进行谐波剔除

    print("----------完成过滤谐波----------")
    # print(amp_new)
    arr1 = np.array(amp_last)
    arr2 = np.array(amp_new)
    amp_change = arr1 - arr2

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 或者其他你喜欢的中文字体
    plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

    ax1.plot(freq_axis, amp_last)
    ax1.set_xlabel("频率")
    ax1.set_ylabel("未滤波的幅度")

    ax2.plot(freq_axis, amp_new)
    ax2.set_xlabel("频率")
    ax2.set_ylabel("滤波后的幅度")
    plt.show()

    return freq_axis, amp_new, sample_rate


# if __name__ == '__main__':
#     path = "../音乐/花海示例音乐.wav"
#     a, b, _ = filter_harmonic(path)

